#include <traj_utils/planning_visualization.h>

using std::cout;
using std::endl;
namespace ego_planner {
Visualization::Visualization(ros::NodeHandle &nh)
{
  node = nh;

  goal_point_pub = nh.advertise<visualization_msgs::Marker>("goal_point", 2);
  global_list_pub = nh.advertise<visualization_msgs::Marker>("global_list", 2);
  init_list_pub = nh.advertise<visualization_msgs::Marker>("init_list", 2);
  optimal_list_pub = nh.advertise<visualization_msgs::Marker>("optimal_list", 2);
  a_star_list_pub = nh.advertise<visualization_msgs::Marker>("a_star_list", 20);

  debug_poly_pub = nh.advertise<visualization_msgs::Marker>("debug_poly", 2);
  debug_inter_pub = nh.advertise<visualization_msgs::Marker>("debug_inter_point", 2);
  debug_yaw_arrow_pub = nh.advertise<visualization_msgs::MarkerArray>("debug_yaw_arrow", 2);
}

// real ids used: {id, id+1000}
void Visualization::displayLine(
    ros::Publisher &pub,
    const vector<Eigen::Vector3d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  visualization_msgs::Marker line_strip;
  line_strip.header.frame_id = "world";
  line_strip.header.stamp = ros::Time::now();
  line_strip.type = visualization_msgs::Marker::LINE_STRIP;
  line_strip.action = visualization_msgs::Marker::ADD;
  line_strip.id = id;

  line_strip.pose.orientation.w = 1.0;
  line_strip.color.r = color(0);
  line_strip.color.g = color(1);
  line_strip.color.b = color(2);
  line_strip.color.a = color(3) > 1e-5 ? color(3) : 1.0;
  line_strip.scale.x = scale / 2;
  geometry_msgs::Point pt;
  for (int i = 0; i < int(list.size()); i++)
  {
    pt.x = list[i](0);
    pt.y = list[i](1);
    pt.z = list[i](2);
    line_strip.points.push_back(pt);
  }
  pub.publish(line_strip);
}

void Visualization::displaySphereList(
    ros::Publisher &pub,
    const vector<Eigen::Vector3d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  if (pub.getNumSubscribers() == 0) return;
  if (list.empty()) return;

  visualization_msgs::Marker sphere;
  sphere.header.frame_id = "world";
  sphere.header.stamp = ros::Time::now();
  sphere.type = visualization_msgs::Marker::SPHERE_LIST;
  sphere.action = visualization_msgs::Marker::ADD;
  sphere.id = id;

  sphere.lifetime = ros::Duration(1.5);

  sphere.pose.orientation.w = 1.0;
  sphere.color.r = color(0);
  sphere.color.g = color(1);
  sphere.color.b = color(2);
  sphere.color.a = color(3) > 1e-5 ? color(3) : 1.0;
  sphere.scale.x = scale;
  sphere.scale.y = scale;
  sphere.scale.z = scale;
  geometry_msgs::Point pt;
  for (int i = 0; i < int(list.size()); i++)
  {
    pt.x = list[i](0);
    pt.y = list[i](1);
    pt.z = list[i](2);
    sphere.points.push_back(pt);
  }
  pub.publish(sphere);
}

void Visualization::displaySphereAndLineList(
    ros::Publisher &pub,
    const vector<Eigen::Vector3d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  visualization_msgs::Marker sphere, line_strip;
  sphere.header.frame_id = line_strip.header.frame_id = "world";
  sphere.header.stamp = line_strip.header.stamp = ros::Time::now();
  sphere.type = visualization_msgs::Marker::SPHERE_LIST;
  line_strip.type = visualization_msgs::Marker::LINE_STRIP;
  sphere.action = line_strip.action = visualization_msgs::Marker::ADD;
  sphere.id = id;
  line_strip.id = id + 1000;

  sphere.pose.orientation.w = line_strip.pose.orientation.w = 1.0;
  sphere.color.r = line_strip.color.r = color(0);
  sphere.color.g = line_strip.color.g = color(1);
  sphere.color.b = line_strip.color.b = color(2);
  sphere.color.a = line_strip.color.a = color(3) > 1e-5 ? color(3) : 1.0;
  sphere.scale.x = scale;
  sphere.scale.y = scale;
  sphere.scale.z = scale;
  line_strip.scale.x = scale / 2;
  geometry_msgs::Point pt;
  for (int i = 0; i < int(list.size()); i++)
  {
    pt.x = list[i](0);
    pt.y = list[i](1);
    pt.z = list[i](2);
    sphere.points.push_back(pt);
    line_strip.points.push_back(pt);
  }
  pub.publish(sphere);
  pub.publish(line_strip);
}

// real ids used: {id, id+1}
void Visualization::generatePathDisplayArray(
    visualization_msgs::MarkerArray &array,
    const vector<Eigen::Vector3d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  visualization_msgs::Marker sphere, line_strip;
  sphere.header.frame_id = line_strip.header.frame_id = "world";
  sphere.header.stamp = line_strip.header.stamp = ros::Time::now();
  sphere.type = visualization_msgs::Marker::SPHERE_LIST;
  line_strip.type = visualization_msgs::Marker::LINE_STRIP;
  sphere.action = line_strip.action = visualization_msgs::Marker::ADD;
  sphere.id = id;
  line_strip.id = id + 1;

  sphere.pose.orientation.w = line_strip.pose.orientation.w = 1.0;
  sphere.color.r = line_strip.color.r = color(0);
  sphere.color.g = line_strip.color.g = color(1);
  sphere.color.b = line_strip.color.b = color(2);
  sphere.color.a = line_strip.color.a = color(3) > 1e-5 ? color(3) : 1.0;
  sphere.scale.x = scale;
  sphere.scale.y = scale;
  sphere.scale.z = scale;
  line_strip.scale.x = scale / 3;
  geometry_msgs::Point pt;
  for (int i = 0; i < int(list.size()); i++)
  {
    pt.x = list[i](0);
    pt.y = list[i](1);
    pt.z = list[i](2);
    sphere.points.push_back(pt);
    line_strip.points.push_back(pt);
  }
  array.markers.push_back(sphere);
  array.markers.push_back(line_strip);
}

// real ids used: {1000*id ~ (arrow nums)+1000*id}
// real ids used: {id ~ (arrow nums)+id}
void Visualization::generateArrowDisplayArray(
    visualization_msgs::MarkerArray &array,
    const vector<Eigen::Vector3d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  // clear
  visualization_msgs::Marker clear_all;
  clear_all.action = visualization_msgs::Marker::DELETEALL;
  clear_all.id = 0;
  array.markers.push_back(clear_all);

  visualization_msgs::Marker arrow;
  arrow.header.frame_id = "world";
  arrow.header.stamp = ros::Time::now();
  arrow.type = visualization_msgs::Marker::ARROW;
  arrow.action = visualization_msgs::Marker::ADD;

  arrow.color.r = color(0);
  arrow.color.g = color(1);
  arrow.color.b = color(2);
  arrow.color.a = color(3) > 1e-5 ? color(3) : 1.0;
  arrow.scale.x = scale;
  arrow.scale.y = 2 * scale;
  arrow.scale.z = 2 * scale;
  arrow.pose.orientation.w = 1.0;

  geometry_msgs::Point start, end;
  for (int i = 0; i < int(list.size() / 2); i++)
  {
    start.x = list[2 * i](0);
    start.y = list[2 * i](1);
    start.z = list[2 * i](2);
    end.x = list[2 * i + 1](0);
    end.y = list[2 * i + 1](1);
    end.z = list[2 * i + 1](2);
    arrow.points.clear();
    arrow.points.push_back(start);
    arrow.points.push_back(end);
    // arrow.id = i + id * 1000;
    arrow.id++;

    array.markers.push_back(arrow);
  }
}

void Visualization::displayGoalPoint(
    Eigen::Vector3d goal_point,
    Eigen::Vector4d color,
    const double scale,
    int id)
{
  visualization_msgs::Marker sphere;
  sphere.header.frame_id = "world";
  sphere.header.stamp = ros::Time::now();
  sphere.type = visualization_msgs::Marker::SPHERE;
  sphere.action = visualization_msgs::Marker::ADD;
  sphere.id = id;

  sphere.pose.orientation.w = 1.0;
  sphere.color.r = color(0);
  sphere.color.g = color(1);
  sphere.color.b = color(2);
  sphere.color.a = color(3);
  sphere.scale.x = scale;
  sphere.scale.y = scale;
  sphere.scale.z = scale;
  sphere.pose.position.x = goal_point(0);
  sphere.pose.position.y = goal_point(1);
  sphere.pose.position.z = goal_point(2);

  goal_point_pub.publish(sphere);
}

void Visualization::displayGlobalPathList(
    vector<Eigen::Vector3d> init_pts,
    const double scale,
    int id)
{

  if (global_list_pub.getNumSubscribers() == 0)
  {
    return;
  }

  Eigen::Vector4d color(0, 0.5, 0.5, 1);
  displayLine(global_list_pub, init_pts, scale, color, id);
}

void Visualization::displayMultiInitPathList(
    vector<vector<Eigen::Vector3d>> init_trajs,
    const double scale)
{

  if (init_list_pub.getNumSubscribers() == 0)
  {
    return;
  }

  static int last_nums = 0;

  for (int id = 0; id < last_nums; id++)
  {
    Eigen::Vector4d color(0, 0, 0, 0);
    vector<Eigen::Vector3d> blank;
    displayLine(init_list_pub, blank, scale, color, id);
    ros::Duration(0.001).sleep();
  }
  last_nums = 0;

  for (size_t id = 0; id < init_trajs.size(); id++)
  {
    Eigen::Vector4d color(0, 0, 1, 0.7);
    displayLine(init_list_pub, init_trajs[id], scale, color, id);
    ros::Duration(0.001).sleep();
    last_nums++;
  }
}

void Visualization::displayInitPathList(
    vector<Eigen::Vector3d> init_pts,
    const double scale,
    int id)
{

  if (init_list_pub.getNumSubscribers() == 0)
  {
    return;
  }

  Eigen::Vector4d color(0, 0, 1, 1);
  displayLine(init_list_pub, init_pts, scale, color, id);
}

void Visualization::displayOptimalList(
    Eigen::MatrixXd optimal_pts,
    int id)
{

  if (optimal_list_pub.getNumSubscribers() == 0)
  {
    return;
  }

  vector<Eigen::Vector3d> list;
  for (int i = 0; i < optimal_pts.cols(); i++)
  {
    Eigen::Vector3d pt = optimal_pts.col(i).transpose();
    list.push_back(pt);
  }
  Eigen::Vector4d color(1, 0, 1, 1);
  displaySphereList(optimal_list_pub, list, 0.15, color, id);
}

void Visualization::displayAStarList(
    std::vector<std::vector<Eigen::Vector3d>> a_star_paths,
    int id /* = Eigen::Vector4d(0.5,0.5,0,1)*/)
{

  if (a_star_list_pub.getNumSubscribers() == 0)
  {
    return;
  }

  int i = 0;
  vector<Eigen::Vector3d> list;
  // make the A star pathes different every time.
  Eigen::Vector4d color{0.5 + ((double)rand() / RAND_MAX / 2),
                        0.5 + ((double)rand() / RAND_MAX / 2),
                        0, 1};
  // Eigen::Vector4d color(0.5,0.5,0,1);
  double scale = 0.05 + (double)rand() / RAND_MAX / 10;

  for (auto block : a_star_paths)
  {
    list.clear();
    for (auto pt : block)
    {
      list.push_back(pt);
    }
    // real ids used: [ id ~ id+a_star_paths.size() ]
    displayLine(a_star_list_pub, list, scale, color, id + i);
    i++;
  }
}

void Visualization::displayArrowList(
    ros::Publisher &pub,
    const vector<Eigen::Vector3d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  visualization_msgs::MarkerArray array;

  // new
  generateArrowDisplayArray(array, list, scale, color, id);

  pub.publish(array);
}

void Visualization::displayPolygon(
    const vector<Eigen::Vector2d> &list,
    double scale,
    int id)
{
  if (list.empty()) return;

  Eigen::Vector4d color1(0, 1, 0, 1);
  pubPolygon(list, scale, color1, id);
}

void Visualization::pubPolygon(
    const vector<Eigen::Vector2d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  if (debug_poly_pub.getNumSubscribers() <= 0) return;
  if (list.empty()) return;

  visualization_msgs::Marker sphere, line_strip;
  sphere.header.frame_id = line_strip.header.frame_id = "world";
  sphere.header.stamp = line_strip.header.stamp = ros::Time::now();
  sphere.type = visualization_msgs::Marker::SPHERE_LIST;
  line_strip.type = visualization_msgs::Marker::LINE_STRIP;
  sphere.id = id;
  line_strip.id = id + 1000;

  sphere.action = visualization_msgs::Marker::DELETE;
  debug_poly_pub.publish(sphere);

  sphere.action = line_strip.action = visualization_msgs::Marker::ADD;

  sphere.pose.orientation.w = line_strip.pose.orientation.w = 1.0;
  sphere.color.r = line_strip.color.r = color(0);
  sphere.color.g = line_strip.color.g = color(1);
  sphere.color.b = line_strip.color.b = color(2);
  sphere.color.a = line_strip.color.a = color(3) > 1e-5 ? color(3) : 1.0;
  sphere.scale.x = scale;
  sphere.scale.y = scale;
  sphere.scale.z = scale;
  line_strip.scale.x = scale / 2;

  geometry_msgs::Point pt;
  for (size_t i = 0; i < list.size(); i++)
  {
    pt.x = list[i](0);
    pt.y = list[i](1);
    pt.z = 0.5;
    sphere.points.push_back(pt);
    line_strip.points.push_back(pt);
  }
  pt.x = list[0](0);
  pt.y = list[0](1);
  pt.z = 0.5;
  line_strip.points.push_back(pt);

  debug_poly_pub.publish(sphere);
  debug_poly_pub.publish(line_strip);
}

void Visualization::displayInterPoint(
    const vector<Eigen::Vector2d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  if (debug_inter_pub.getNumSubscribers() == 0) return;

  visualization_msgs::Marker sphere;
  sphere.header.frame_id = "world";
  sphere.header.stamp = ros::Time::now();
  sphere.type = visualization_msgs::Marker::SPHERE_LIST;
  sphere.id = id;

  // delete
  sphere.action = visualization_msgs::Marker::DELETE;
  debug_inter_pub.publish(sphere);

  sphere.action = visualization_msgs::Marker::ADD;

  sphere.pose.orientation.w = 1.0;
  sphere.color.r = color(0);
  sphere.color.g = color(1);
  sphere.color.b = color(2);
  sphere.color.a = color(3) > 1e-5 ? color(3) : 1.0;
  sphere.scale.x = scale;
  sphere.scale.y = scale;
  sphere.scale.z = scale;
  geometry_msgs::Point pt;
  for (int i = 0; i < int(list.size()); i++)
  {
    pt.x = list[i](0);
    pt.y = list[i](1);
    pt.z = 0.5;
    sphere.points.push_back(pt);
  }

  debug_inter_pub.publish(sphere);
}

void Visualization::displayYawArrowList(
    const vector<Eigen::Vector3d> &list,
    double scale,
    Eigen::Vector4d color,
    int id)
{
  if (debug_yaw_arrow_pub.getNumSubscribers() == 0) return;

  visualization_msgs::MarkerArray array;

  // new
  generateArrowDisplayArray(array, list, scale, color, id);

  debug_yaw_arrow_pub.publish(array);
}

} // namespace ego_planner